/* 
 * This code is part of project 
 * "Sårbarheder i og trusler mod trådløse netværk med fokus på WEP."
 * Projects in IT-Security and Cryptography - Department of Computer Science - Aarhus University 
 * April - June 2015
 * 
 * Jesper René Frandsen 
 * George Catalin Hategan 
 * Torben Krogh Johansen
 */
package krypto.projekt;

/**
 * This is a implementation of RC4 algorithm.
 *
 * "NormaL" RC4 operation:
 *  - It can be used either to generate a stream to use your self, or use the built in XOR function.
 *
 * "Hacker mode" RC4 operation:
 *  - Used to control PRGA sequence. Then check if the internal state is in a resolved state.
 *    If the state is resolved, it is possible to guess one byte of the key.
 *  
 * RC4 is split into two phases: 
 * 	- KSA - Key Scheduling Algorithm.
 * 	- PRGA - Pseudo Random Generator Algorithm.
 *
 * In the KSA phase, the internal state S is created and scrambled.
 * In the PRGA phase a stream is generated. 
 * This stream can be used to encrypt or decrypt by XOR'ing.
 */
public class RC4Impl {
	final static int SIZE = 256;

	int[] S = new int[SIZE];			
	int[] K;

	int iUsedByPRGASequence = 0;
	int jUsedByPRGASequence = 0;

	int keyGuessJValue = 0;
	int keyGuess = 0;

	//Used to indicate if KSA sequence is completed. If not, the PRGA sequence should not be started.
	boolean isKSACompleted = false;

	public RC4Impl() {
		super();
	}

	/**
	 * "Normal mode". 
	 * Validate key, initialize internal state S and initialize variables used in PRGA phase
	 */
	public void doKSA(int[] K) throws Exception {
		validateKey(K); 
		this.K = K;

		initAndScramble(SIZE);

		//These variables are used in PRGA phase
		iUsedByPRGASequence = 0;
		jUsedByPRGASequence = 0;

		isKSACompleted = true;
	}

	/**
	 * "Hacker mode"
	 * Validate key, initialize internal state S for a certain amount of cycles
	 */
	public boolean isResolved(int[] K, int l, int x0) throws Exception {
		isKSACompleted = false;

		validateKey(K); 
		this.K = K;

		initAndScramble(l);
		boolean test1 = false;
		boolean test2 = false;
		boolean test3 = false;
		boolean test4 = false;

		if (S[1] < l) { 	
			test1 = true;
		} 

		if ((S[1] + S[S[1]]) == l) { 
			test2 = true;
		}

		if (getInvPerm(x0) != 1) { 
			test3 = true;
		}

		if (getInvPerm(x0) != S[1]) { 
			test4 = true;
		}

		if (test1 && test2 && test3 && test4) {
			keyGuess = mod((getInvPerm(x0) - keyGuessJValue - S[l]),SIZE);
			return true;
		} else {
			return false;
		}
	}
 
	/**
	 * Used to return guess. Only makes sense to use if state is resolved.
	 */
	public int getKeyGuess() {
		return keyGuess;
	}

	/**
	 * Do a lookup in the internal state S to get offset containing a specific value
	 */
	public int getInvPerm(int value) {
		int pos = 0;

		for (int i=0; i<S.length; i++) {
			if (S[i] == value) {
				pos = i;
				break;
			}
		}

		return pos;
	}

	/**
	 * Internal state S is initialized with values from zero and up.
	 * Afterwards key is used to scramble the internal state.
	 * 
	 * The method takes a integer as input which is used to control the
	 * cramble part. Used to limit scramble sequence in "hacker mode"
	 */
	public void initAndScramble(int scrambleCount) {
		for (int i=0; i<SIZE; i++) {
			S[i] = i;
		}

		int swapTemp;
		int j = 0;

		//Scrambling / generate a random permutation
		for (int i=0; i<scrambleCount; i++) {
			j = mod((j + S[i] + K[mod(i,K.length)]),SIZE);

			swapTemp = S[i];
			S[i] = S[j];
			S[j] = swapTemp;
		}

		//This value is needed if we are trying to guess code / cracking
		keyGuessJValue = j;
	}

	/**
	 * Key is validated and exception is thrown if the key is not ok.
	 */
	public void validateKey(int[] K) throws Exception {
		//validate key not null
		if (K == null) {
			throw new Exception("Key must not be null");
		}

		//validate key length
		if (K.length > 256) {
			throw new Exception("Key length may not be larger than 256");
		}

		//validate content
		for (int tempKey : K ) {
			if (tempKey > 256) {
				throw new Exception("Key may not contain values larger than 256");
			}
		}
	}

	/**
	 * This is the PRGA part of RC4.
	 * Gives one value at a time. So in order to get a stream,
	 * you call the method again and again.
	 */
	public int getNextFromPRGA() throws Exception {
		int i = iUsedByPRGASequence;
		int j = jUsedByPRGASequence;

		if (!isKSACompleted) {
			throw new Exception("KSA has do be completed before PRGA ");
		}

		i = mod((i + 1),SIZE);
		j = mod((j + S[i]),SIZE);

		int swapTemp = S[i];
		S[i] = S[j];
		S[j] = swapTemp;

		int index = mod((S[i] + S[j]),SIZE);

		iUsedByPRGASequence = i;
		jUsedByPRGASequence = j;

		return S[index];
	}

	/**
	 * This can be used to encrypt or decrypt using the getNextFromPRGA method in this class
	 */
	public int[] XORIntArray(int[] input) throws Exception{
		int[] output = new int[input.length];

		for (int i=0; i<input.length; i++) {
			//XOR Using RC4 stream
			output[i] = input[i]^getNextFromPRGA();
		}

		return output;
	}

	/**
	 * Modulus operation in Java could be much better.
	 * This one is used in stead. It does not return negative values. 
	 */
	public int mod(int i, int j) {		
		int temp = i%j;

		if (temp < 0)
		{
			temp = temp + j;
		}

		return temp;
	}
}
